from efficientdet_model.bifpn import BIFPN
from efficientdet_model.efficientnet import EfficientNet
import torch.nn as nn
import torch
class Efficient_Bifpn(nn.Module):
    def __init__(self,effic_name = 'efficientnet-b0'):
        super(Efficient_Bifpn, self).__init__()
        self.backbone = EfficientNet.from_name(effic_name)
        self.bifpn = BIFPN(in_channels=self.backbone.get_list_features()[-5:-2],
                      out_channels=88,
                      stack=3,
                      num_outs=5)
    def forward(self,input):
        backbone_features = self.backbone(input)
        bifpn_features    = self.bifpn(backbone_features[-5:-2])
        return bifpn_features
    
if __name__ == '__main__':
    input = torch.randn(1,3,512,512)
    module = Efficient_Bifpn()
    features = module(input)
    print([i.shape for i in features])
    